import torch
import torch.nn as nn
import numpy as np
from utils.utils import *
from utils.bbox import *


YoloLossHyp = {'iou': 3.54,
               'cls': 1.181,  # cls loss fraction
               'obj': 1.509,  # obj conf loss fraction
               'obj_pw': 1.0,  # obj BCELoss positive_weight
               'cls_pw': 1.0,   # cls BCELoss positive_weight
               }


def _build_target(targets, model):
    """
    Build targets for compute_loss()
    :param targets:     GroundTruth, [batch_id, class_id, cx, cy, w, h], n x 6
    :param model:       Training model, contains hyperparameters and yolo layers
    :return:
    """
    assert(targets is not None)
    FloatTensor = torch.cuda.FloatTensor if targets.is_cuda else torch.FloatTensor
    IntTensor = torch.cuda.IntTensor if targets.is_cuda else torch.IntTensor
    nt = targets.shape[0]
    nc = model.yolo_layers[0].nc
    # 遍历所有YOLO层的anchor，并计算每个层中GroundTruth BBox与该层anchor的宽高IOU
    anchors_all, ious_all = [], []
    ngy_all, ngx_all = [], []
    for i, yolo in enumerate(model.yolo_layers):
        anchors = yolo.scaled_anchors
        ngy = yolo.ngy
        ngx = yolo.ngx

        gain = torch.ones([nt, 6], device=targets.device)
        gain[:, 2:] *= FloatTensor([ngx, ngy, ngx, ngy])
        gts = targets * gain
        ious = torch.stack([bbox_wh_iou(anchor, gts[:, 4:6]) for anchor in anchors])
        ious_all.append(ious)
        anchors_all.append(anchors)
        ngy_all.append(ngy)
        ngx_all.append(ngx)

    anchors = torch.cat(anchors_all, 0) # 合并所有YOLO层的anchor
    ious = torch.cat(ious_all, 0)       # 合并所有GroundTruth BBox与所有anchor的IOU值
    ngx = IntTensor(ngx_all)            # 合并所有YOLO层的x方向网格数目
    ngy = IntTensor(ngy_all)            # 合并所有YOLO层的y方向网格数目
    # 找出GroundTruth BBox与所有anchor的最大IOU, 记录最佳匹配的anchor索引
    best_iou, best_n = ious.max(dim=0)
    ny = best_n // len(model.yolo_layers)   # 计算最优anchor所属的YOLO层索引
    na = best_n % len(model.yolo_layers)    # 计算最优anchor所属的YOLO层中的子索引
    gain = torch.ones([nt, 6], device=targets.device)
    gain[:, 2:] *= torch.stack([ngx[ny], ngy[ny], ngx[ny], ngy[ny]]).t()
    gts = targets * gain
    b, labels = gts[:, :2].long().t()   # b: 批次索引, labels: 类别索引
    gx, gy, gw, gh = gts[:, 2:].t()     # GT BBox在对应YOLO层中的相对尺寸
    gi = gx.long()  # gi : [nt], 向下取整，获得所属cell的横坐标
    gj = gy.long()  # gj : [nt], 向下取整，获得所属cell的纵坐标
    # 转换GT BBox的格式与检测模型输出格式一致
    # tx = gx
    # ty = gy
    # tw = torch.log(gw / anchors[best_n, 0] + 1e-16)
    # th = torch.log(gh / anchors[best_n, 1] + 1e-16)
    # tbox = torch.stack([tx, ty, tw, th]).t()
    tbox = gts[:, 2:]
    indices = torch.stack([b, na, gj, gi]).t()
    yls = ny;

    tcls = torch.zeros([nt, nc], device=targets.device)
    tcls[range(nt), labels] = 1.0
    return tbox, tcls, indices, yls


def compute_loss(pred, targets, model, lossHyp = YoloLossHyp):
    FloatTensor = torch.cuda.FloatTensor if targets.is_cuda else torch.FloatTensor
    tbox, tcls, indices, yls = _build_target(targets, model)
    BCEcls = nn.BCEWithLogitsLoss(pos_weight=FloatTensor([lossHyp["cls_pw"]]), reduction="mean")
    BCEobj = nn.BCEWithLogitsLoss(pos_weight=FloatTensor([lossHyp["obj_pw"]]), reduction="mean")

    loss_ciou, loss_cls, loss_obj = FloatTensor([0]), FloatTensor([0]), FloatTensor([0])
    # 定义损失度量， 6 x ny, 每行分别代表 total_loss, ciou_loss, obj_loss, cls_loss, grid_x, grid_y
    lossMetric = torch.zeros([6, len(pred)]).cpu()
    for i, predict in enumerate(pred):
        ngx = model.yolo_layers[i].ngx
        ngy = model.yolo_layers[i].ngy
        nc = model.yolo_layers[i].nc
        aw = model.yolo_layers[i].anchor_w
        ah = model.yolo_layers[i].anchor_h

        # 填充度量变量
        lossMetric[4, i] = ngx
        lossMetric[5, i] = ngy

        # 还原预测坐标到当前尺度层上
        pbox = torch.zeros_like(predict[..., 0:4], device=predict.device)
        px = torch.sigmoid(predict[..., 0]) + ngx
        py = torch.sigmoid(predict[..., 1]) + ngy
        pw = torch.exp(predict[..., 2]) * aw
        ph = torch.exp(predict[..., 3]) * ah
        pbox[..., 0] = px
        pbox[..., 1] = py
        pbox[..., 2] = pw
        pbox[..., 3] = ph
        # 获取预测的置信度
        pconf = torch.sigmoid(predict[..., 4])
        pcls = torch.sigmoid(predict[..., 5:])

        # 抓取属于当前YOLO层的索引序列
        idx = yls == i
        index = indices[idx]
        bi, ai, gyi, gxi = index.t()

        if idx.sum() == 0:  # 若没有分配到当前层的GroundTruth

            # ciou损失取0
            lciou = 0
            loss_ciou += lciou
            lossMetric[1, i] = lciou
            # 计算Obj置信度损失
            gconf = torch.zeros_like(pconf)  # 无目标，GroundTruth 置信度全为0
            lobj = BCEobj(pconf, gconf) * lossHyp["obj"]  # 整个obj置信度张量全部参与计算，包含目标与背景
            loss_obj += lobj
            lossMetric[2, i] = lobj.detach().cpu()
            # Cls置信度损失取0
            lcls = 0
            loss_cls += lcls
            lossMetric[3, i] = lcls
        else:
            # 抓取分配到当前YOLO层上的 GroundTruth BBox
            gbox = tbox[idx]
            # 抓取分配到当前YOLO层上的 GroundTruth 类别
            gcls = tcls[idx]
            # 使用One-hot编码当前YOLO层上的 GroundTruth Obj置信度
            gconf = torch.zeros_like(pconf)
            gconf[bi, ai, gyi, gxi] = 1.0   # one-hot encode
            # 计算CIoU损失
            cious = bbox_ciou(pbox[bi, ai, gyi, gxi, :], gbox, x1y1x2y2=False)
            lciou = (1.0 - cious).mean() * lossHyp["iou"]
            loss_ciou += lciou
            lossMetric[1, i] = lciou.detach().cpu()
            # 计算Obj置信度损失
            lobj = BCEobj(pconf, gconf) * lossHyp["obj"] # 整个obj置信度张量全部参与计算，包含目标与背景
            loss_obj += lobj
            lossMetric[2, i] = lobj.detach().cpu()
            # 计算Cls置信度损失
            lcls = BCEcls(pcls[bi, ai, gyi, gxi, :], gcls) * lossHyp["cls"]
            loss_cls += lcls
            lossMetric[3, i] = lcls.detach().cpu()
        # 计算该YOLO层的总损失
        lossMetric[0, i] = lossMetric[1:4, i].sum()

    loss = (loss_ciou + loss_cls + loss_obj)

    return loss, lossMetric
